package com.bigkai.dakaiaiagent.chatmemory;

import com.esotericsoftware.kryo.Kryo;
import com.esotericsoftware.kryo.io.Input;
import com.esotericsoftware.kryo.io.Output;
import org.objenesis.strategy.StdInstantiatorStrategy;
import org.springframework.ai.chat.memory.ChatMemory;
import org.springframework.ai.chat.messages.Message;

import java.io.File;
import java.io.FileInputStream;
import java.io.FileOutputStream;
import java.io.IOException;
import java.util.ArrayList;
import java.util.List;

/**
 * 基于文件持久化的对话记忆
 */
public class FileBasedChatMemory implements ChatMemory {

  private final String BASE_DIR;
  private static final Kryo kryo = new Kryo();

  static {
    kryo.setRegistrationRequired(false);
    // 设置实例化策略
    kryo.setInstantiatorStrategy(new StdInstantiatorStrategy());
  }

  // 构造对象时，指定文件保存目录
  public FileBasedChatMemory(String dir) {
    this.BASE_DIR = dir;
    File baseDir = new File(dir);
    if (!baseDir.exists()) {
      baseDir.mkdirs();
    }
  }

  @Override
  public void add(String conversationId, List<Message> messages) {
    List<Message> conversationMessages = getOrCreateConversation(conversationId);
    conversationMessages.addAll(messages);
    saveConversation(conversationId, conversationMessages);
  }

  @Override
  public List<Message> get(String conversationId) {
    return getOrCreateConversation(conversationId);
  }


  @Override
  public void clear(String conversationId) {
    File file = getConversationFile(conversationId);
    if (file.exists()) {
      file.delete();
    }
  }

  private List<Message> getOrCreateConversation(String conversationId) {
    File file = getConversationFile(conversationId);
    List<Message> messages = new ArrayList<>();
    if (file.exists()) {
      try (Input input = new Input(new FileInputStream(file))) {
        messages = kryo.readObject(input, ArrayList.class);
      } catch (IOException e) {
        e.printStackTrace();
      }
    }
    return messages;
  }

  private void saveConversation(String conversationId, List<Message> messages) {
    File file = getConversationFile(conversationId);
    try (Output output = new Output(new FileOutputStream(file))) {
      kryo.writeObject(output, messages);
    } catch (IOException e) {
      e.printStackTrace();
    }
  }

  private File getConversationFile(String conversationId) {
    return new File(BASE_DIR, conversationId + ".kryo");
  }
}